package com.xes.hardwritingrecognition;

import android.content.Context;
import android.content.res.AssetManager;

import com.xes.hardwritingrecognition.pytorch.PytorchClassifier;
import com.xes.hardwritingrecognition.tensorflow.TensorFlowImageClassifier;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

public class MainExecutor {
    public enum ClassifierType {
        TensorFlow, PyTorch
    }
    public static ClassifierType classifierType = ClassifierType.PyTorch;
    private static final int INPUT_SIZE = 28;
    private static final String INPUT_NAME = "input";
    private static final String OUTPUT_NAME = "output";
    private static final String LABEL_FILE = "graph_label_strings.txt";

    private Executor executor;
    private Classifier classifier;

    public MainExecutor(Context context, InputStream modelFileInputStream, String assetsName, boolean isAssets) {
        executor = Executors.newSingleThreadExecutor();

        switch (classifierType) {
            case TensorFlow: {
                initTensorFlowClassifier(context, modelFileInputStream, assetsName, isAssets);
                break;
            }
            case PyTorch: {
                initPyTorchClassifier(context, modelFileInputStream, assetsName, isAssets);
                break;
            }
        }
    }

    private void initTensorFlowClassifier(final Context context, final InputStream modelInputStream, final String assetsName, final boolean isAssets) {
        executor.execute(new Runnable() {
            @Override
            public void run() {
                try {
                    InputStream is = modelInputStream;
                    if (isAssets) {
                        is = context.getAssets().open(assetsName);
                    }
                    classifier = TensorFlowImageClassifier.create(
                            is,
                            context.getAssets().open(LABEL_FILE),
                            INPUT_SIZE,
                            INPUT_NAME,
                            OUTPUT_NAME);
                } catch (final Exception e) {
                    throw new RuntimeException("Error initializing TensorFlow!", e);
                }
            }
        });
    }

    private void initPyTorchClassifier(final Context context, final InputStream modelInputStream, final String assetsName, final boolean isAssets) {
        executor.execute(new Runnable() {
            @Override
            public void run() {
                try {
                    if (isAssets) {
                        classifier = PytorchClassifier.create(context, assetsName, INPUT_SIZE);
                    } else {
                        classifier = PytorchClassifier.create(context, modelInputStream, INPUT_SIZE);
                    }
                } catch (Exception e) {
                    throw new RuntimeException("Error initializing PyTorch!", e);
                }
            }
        });
    }

    //获得预测结果
    public List<Classifier.Recognition> getResult(float[] data) {
        if (classifier == null) {
            return new ArrayList<>();
        }
        return classifier.recognizeImage(data);
    }

    public void enableLogging(boolean enable) {
        if (classifier != null) {
            classifier.enableStatLogging(enable);
        }
    }

    public void closeClassifier() {
        executor.execute(new Runnable() {
            @Override
            public void run() {
                classifier.close();
            }
        });
    }

}


